import numpy as np


class PosTrains:
    '''
        sucker_height: 吸盘高度
        tetris_height: 方块高度
        block_size: 盘子方格尺寸
        left_down_center: 盘子左下角方块的中心坐标
    '''
    sucker_height = 0.009
    tetris_height = 0.009
    block_size = 0.0185

    plate_center = np.array([-0.5, 0.3, 0.4])
    left_down_center = plate_center + np.array([block_size*6.5, -block_size*4.5, 0])

    @classmethod
    def TrainSport(cls, x:int, y:int) -> np.ndarray[float]:
        '''
            将盘子坐标（二维矩阵形式）转换为仿真环境坐标。
            以盘子左下角为原点，向右为x正半轴，向上为y正半轴。
            return: (x, y, z)
        '''
        off = np.array([- y * cls.block_size,   # x轴偏移量
                        x * cls.block_size,     # y轴偏移量
                        cls.tetris_height])     # z轴偏移量, 方块的高度
        real_pos = cls.left_down_center + off
        return real_pos
    
    @classmethod
    def TrainsPortBlockPos(cls, block_name:str, position:np.ndarray) -> np.ndarray[float]:
        '''
            转换方块坐标，将中心坐标转换到方格上，以便吸附。
            block_name: color + index + "/", 如"orange0/"
            return: (x, y, z)
        '''
        if block_name[:4] == "blue":
            return position + np.array([cls.block_size / 2, 0, cls.tetris_height])
        elif block_name[:4] == "brow":
            return position + np.array([0, 0, cls.tetris_height])
        elif block_name[:4] == "gree":
            return position + np.array([- cls.block_size / 2, 0, cls.tetris_height])
        elif block_name[:4] == "oran":
            return position + np.array([cls.block_size / 2, cls.block_size / 2, cls.tetris_height])
        elif block_name[:4] == "purp":
            return position + np.array([- cls.block_size / 2, 0, cls.tetris_height])
        elif block_name[:3] == "red":
            return position + np.array([- cls.block_size / 2, 0, cls.tetris_height])
        elif block_name[:4] == "yell":
            return position + np.array([cls.block_size / 2, 0, cls.tetris_height])